from dataclasses import dataclass
from typing import Sequence, Union

import numpy as np
from gymnasium import spaces
from ..base_controller import BaseController, ControllerConfig
import sapien.core as sapien
import time
class PDBasePosController():
    def __init__(self,config: "PDBasePosControllerConfig"):
        self.config = config
        self.dones = False
        self.dist_qpos = 10000
        self.reset()
        self.set_pose()
        self.delay()
        # if(self.dones==True)

    def initialize_action_space(self):
        low,high = np.array(self.config.lower),np.array(self.config.upper)
        self.action_space = spaces.Box(low, high, dtype=np.float32)
    def set_drive_property(self):
        self.active_joints = self.config.articulation.get_active_joints()
        for joint_idx,joint in enumerate(self.active_joints):
            if(joint_idx<4):
                joint.set_drive_property(self.config.stiffness, self.config.damping)
            # else:
            #     joint.set_drive_property(self.config.stiffness, self.config.damping)
    def reset(self):
        self.initialize_action_space()
        self.set_drive_property()
        self.start_qpos = self.config.articulation.get_qpos()
    def set_drive_targets(self,targets):
        for i, joint in enumerate(self.active_joints):
            joint.set_drive_target(targets[i])
    def set_pose(self):
        if self.config.use_delta:
            self.start_qpos[0:4]+=self.config.target
            self.target_qpos = self.start_qpos
            self.set_drive_targets(self.target_qpos)
        else:
            dist_qpos = np.linalg.norm(self.config.target-self.start_qpos[0:4])
            # print('dist_qpos', dist_qpos)
            self.dist_qpos = dist_qpos
            if dist_qpos <self.config.done_thresh:
                self.dones = True
                return
            self.start_qpos[0:4]=self.config.target
            if(self.config.with_grasp):
                self.start_qpos[4:] =self.config.begin_qpos[4:]
                self.start_qpos[7]=-3
                self.start_qpos[16]=-3
            self.target_qpos = self.start_qpos
            self.set_drive_targets(self.target_qpos)
            
            
    def delay(self):
        last_qpos = self.config.articulation.get_qpos()[0]
        # time.sleep(self.config.time_interval)
        while ((self.config.articulation.get_qpos()[0]-last_qpos)>0.01):
            last_qpos = self.config.articulation.get_qpos()[0]
            # time.sleep(self.config.time_interval)
        





        

@dataclass
class PDBasePosControllerConfig():
    lower: Union[None, float, Sequence[float]]
    upper: Union[None, float, Sequence[float]]
    stiffness: Union[float, Sequence[float]]
    damping: Union[float, Sequence[float]]
    target: Union[None,list]
    articulation:sapien.Articulation
    begin_qpos:np.ndarray
    with_grasp:bool = False
    force_limit: Union[float, Sequence[float]] = 1e10
    friction: Union[float, Sequence[float]] = 0.0
    use_delta: bool = False
    done_thresh: float = 0.1
    time_interval: float = 0.01
    # normalize_action: bool = True
    controller_cls = PDBasePosController
